import torch
import torch.nn as nn
import torch.nn.functional as F
from .activate import *


norm_name = {"bn": nn.BatchNorm2d}
activate_name = {
    "relu": nn.ReLU,
    "leaky": nn.LeakyReLU,
    "HardMish": HardMish}


class Convolutional(nn.Module):
    def __init__(self, filters_in, filters_out, kernel_size, stride=1, pad=0, norm=None, activate=None):
        super(Convolutional, self).__init__()

        self.norm = norm
        self.activate = activate
        self.out_channels = filters_out
        self.__conv = nn.Conv2d(in_channels=filters_in, out_channels=filters_out, kernel_size=kernel_size,
                                stride=stride, padding=pad, bias=not norm)
        if norm:
            # assert norm in norm_name.keys()
            if norm == "bn":
                self.__norm = nn.BatchNorm2d(filters_out)

        if activate:
            # assert activate in activate_name.keys()
            if activate == "leaky":
                self.__activate = nn.LeakyReLU(negative_slope=0.1, inplace=True)
            elif activate == "relu":
                self.__activate = nn.ReLU(inplace = True)
            elif activate == "HardMish":
                self.__activate = activate_name[activate](inplace = True)

    def forward(self, x):
        x = self.__conv(x)
        if self.norm:
            x = self.__norm(x)
        if self.activate:
            x = self.__activate(x)
        return x
